#!/bin/bash

source ../../lib/shell/common.sh

# setup to setup gpu test env
# 1. download and nvidia driver: NVIDIA-Linux-x86_64-470.57.02.run
#    if require to add nouveau to module blacklist, you must select
#    to add nouveau driver to blacklist, then run command:
#    cd /boot
#    dracut -f
# 2. install cude library: cuda_11.2.0_460.27.04_linux.run, do not
#    select to install driver
# 3. if you are using the GPU from Iluvatar CoreX, please contact 
#    their application engineers to obtain the SDK installation 
#    package from https://www.iluvatar.com
# 4. reboot system, then check the driver whether is nvidia or bi

function get_gpu_info() {
	echo "device path: $DEVPATH"
	local bdf
	bdf="$(echo "$DEVPATH" | awk -F'/' '{print $NF}')"

	echo "device info:"
	lspci -vvvv -s "$bdf"
}

function burn_nvidia_gpu() {
	which nvidia-smi > /dev/null 2>&1
        if [ $? -ne 0 ]; then
		echo "nvidia-smi tool does not install, skip"
                return
        fi

	local gpu_burn_log=/tmp/gpu_burn.ancert
	local gpu_burn_pid=
	cd gpu_burn
	# require gcc-g++
	make
	if [ -f gpu_burn ]; then
		./gpu_burn 30 > $gpu_burn_log 2>&1 &
		gpu_burn_pid=$!
	fi

	cd -
	echo "show invidi gpu statistics:"
	for ((i=0;i<10;i++)); do
		nvidia-smi
		sleep 3
	done

	echo "wait gpu burn done"
	wait $gpu_burn_pid

	echo "gpu burn data:"
	cat $gpu_burn_log
	rm -rf $gpu_burn_log
	cd gpu_burn
	make clean
	cd -
}
function burn_hygon_gpu() {
        which hy-smi > /dev/null 2>&1
        if [ $? -ne 0 ]; then
                echo "hy-smi tool does not install, skip"
                return
        fi

        local gpu_burn_log=/tmp/gpu_burn.ancert
        local gpu_burn_pid=
        cd vectorAdd_dcu
        # require gcc-g++
        make
        if [ -f vector_add ]; then
                ./vector_add > $gpu_burn_log 2>&1 &
                gpu_burn_pid=$!
        fi

        cd -
        echo "show haiguang gpu statistics:"
        for ((i=0;i<10;i++)); do
                hy-smi
                sleep 3
        done

        echo "wait gpu burn done"
        wait $gpu_burn_pid

        echo "gpu burn data:"
        cat $gpu_burn_log
        rm -rf $gpu_burn_log
        cd vectorAdd_dcu
        make clean
        cd -
}
function burn_denglin_gpu() {
        which dlsmi > /dev/null 2>&1
        if [ $? -ne 0 ]; then
                echo "dlsmi tool does not install, skip"
                return
        fi

        local gpu_burn_log=/tmp/gpu_burn.ancert
        local gpu_burn_pid=
        cd vectorAdd_nvrtc
        # require gcc-g++
        make
        if [ -f vector_add_nvrtc ]; then
                ./vector_add_nvrtc > $gpu_burn_log 2>&1 &
                gpu_burn_pid=$!
        fi

        cd -
        echo "show denglin gpu statistics:"
        for ((i=0;i<10;i++)); do
                dlsmi
                sleep 3
        done

        echo "wait gpu burn done"
        wait $gpu_burn_pid

        echo "gpu burn data:"
        cat $gpu_burn_log
        rm -rf $gpu_burn_log
        cd vectorAdd_nvrtc
        make clean
        cd -
}
function burn_bi_gpu() {
        which ixsmi > /dev/null 2>&1
        if [ $? -ne 0 ]; then
                echo "ixsmi tool does not install, skip"
                return
        fi

        local gpu_burn_log=/tmp/gpu_burn.ancert
        local gpu_burn_pid=
        cd gpu_burn
        res_code=0
        lsmod | grep bi_driver >/dev/null
        if [ $? -eq 0 ];then
                COREX_PATH=${COREX_PATH:-/usr/local/corex}
                clang++ compare.cu -o compare.ll -I${COREX_PATH}/include --cuda-gpu-arch=ivcore10 --cuda-path=${COREX_PATH} --cuda-device-only -S -x cuda || res_code=1
                llc -mcpu=ivcore10 -mtriple=bi-iluvatar-ilurt -show-mc-encoding -filetype=obj compare.ll -o compare.o || res_code=1
                lld -flavor ld.lld --no-undefined --no-warn-missing-entry compare.o -o compare.ptx || res_code=1
                rm compare.ll compare.o
                sed -i '/cuFuncSet/d' gpu_burn-drv.cpp
                sed -i '/cuParamSet/d' gpu_burn-drv.cpp
                sed -i 's/nvidia-smi/ixsmi/g' gpu_burn-drv.cpp
                sed -i 's/.*cuLaunchGridAsync.*/void\* kargs[] = {\&d_Cdata, \&d_faultyElemData, \&d_iters};checkError(cuLaunchKernel(d_function, SIZE\/g_blockSize, SIZE\/g_blockSize, 1, g_blockSize, g_blockSize, 1, 0, 0, kargs, nullptr));/g' gpu_burn-drv.cpp
                clang++ -std=c++11 -I${COREX_PATH}/include -L${COREX_PATH}/lib64 -lcudart -lcuda -lcublas -o gpu_burn ./gpu_burn-drv.cpp || res_code=1
        else
                make &>/dev/null || res_code=1
        fi

        if [ -f gpu_burn ]; then
                ./gpu_burn 30 > $gpu_burn_log 2>&1 &
                gpu_burn_pid=$!
        fi

        cd -
        echo "show bi gpu statistics:"
        for ((i=0;i<10;i++)); do
                ixsmi
                sleep 3
        done

        echo "wait gpu burn done"
        wait $gpu_burn_pid

        echo "gpu burn data:"
        cat $gpu_burn_log
        rm -rf $gpu_burn_log
        cd gpu_burn
        make clean
        cd -
}

function burn_gpu() {
	if echo "$ID_VENDOR_FROM_DATABASE" | grep -i -q NVIDIA; then
		burn_nvidia_gpu
	elif echo "$ID_VENDOR_FROM_DATABASE" | grep -i -q Iluvatar; then
		burn_bi_gpu
        elif echo "$ID_VENDOR_FROM_DATABASE" | grep -i -q DENGLIN; then
                burn_denglin_gpu
	elif echo "$ID_VENDOR_FROM_DATABASE" | grep -i -q -E "haiguang|higon"; then
		burn_hygon_gpu
	else
		echo "$ID_VENDOR_FROM_DATABASE is not supported!"
		test_fail
	fi
}

# # # # # # # # # # # # # # main # # # # # # # # # # # # # #
print_test_info "GPU"
show_module_info "$MODULE"
burn_gpu
test_pass
